/**
 * @author luzeyu & liangqiaowei
 */
package tiei.ajp.huffman.utils;

import com.sun.source.tree.Tree;
import org.junit.Test;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.PriorityQueue;


public class HuffmanTree<T> {
//    public static TreeNode<Character> buildHuffmanTree(int[] frequencies) {
//        PriorityQueue<TreeNode<Character>> queue = new PriorityQueue<>();
//
//        // 初始化：为每个非零频率的字符创建一个叶节点
//        for (char i = 0; i < frequencies.length; i++) {
//            if (frequencies[i] > 0) {
//                queue.add(new TreeNode(frequencies[i], i));
//            }
//        }
//
//        // 不断合并最低频率的两个节点，直到队列中只剩一个节点
//        while (queue.size() > 1) {
//            TreeNode<Character> left = queue.poll();
//            TreeNode<Character> right = queue.poll();
//
//            // 创建一个新节点，频率为两个子节点的和
//            TreeNode<Character> parent = new TreeNode(left.freq + right.freq, left, right);
//            queue.add(parent);
//        }
//
//        // 剩下的唯一节点即为 Huffman 树的根节点
//        return queue.poll();
//    }

    private PriorityQueue<TreeNode<T>> minHeap = null;
    private boolean hasBuilt = false;
    private TreeNode<T> treeRoot = null;

    private static final byte NON_LEAF_MARK = 1;
    private static final byte LEAF_MARK = 0;

    public HuffmanTree() {
        this.minHeap = new PriorityQueue<>();
    }

    /**
     * Deserialize from byte array
     *
     * @param arr byte array
     */
    public HuffmanTree(List<Byte> arr) {
        List<TreeNode> nodes = new ArrayList<>();

        boolean skipNext = false;
        for (int i = 0; i < arr.size(); i++) {
            // Check if current byte is data
            if (skipNext) {
                skipNext = false;
                continue;
            }
            // --- Build tree nodes
            // Current byte indicates an inner node
            if (arr.get(i) == 0x1) nodes.add((new TreeNode<Byte>()));
                // Current byte indicates a leaf node, next byte is data
            else {
                assert arr.get(i) == 0x0;
                assert i < arr.size() - 1;

                nodes.add(new TreeNode(arr.get(i + 1)));
                skipNext = true; // Next byte is data, need to skip
            }
        }
        this.treeRoot = buildTreeFromPreorderRecursive(nodes);
        hasBuilt = true;
    }

    /**
     * Add node to the priority queue
     *
     * @param node
     */
    public void addNode(TreeNode<T> node) {
        assert node.data != null; // Do not allow to add non-leaf node
        assert !hasBuilt; // Do not allow adding to already built tree
        this.minHeap.add(node);
    }

    public TreeNode<T> buildTree() {
        assert !hasBuilt;

        if (minHeap.isEmpty()) {
            return null;
        }

        while (minHeap.size() > 1) {
            TreeNode<T> left = minHeap.poll();
            TreeNode<T> right = minHeap.poll();
            assert right != null;
            TreeNode<T> parent = new TreeNode<T>(left.freq + right.freq, left, right);
            minHeap.add(parent);
        }

        assert minHeap.size() == 1;

        hasBuilt = true;
        treeRoot = minHeap.peek();
        return minHeap.poll();
    }


    private List<TreeNode<T>> nodeList = null;

    private void preorderTraverseRecursive(TreeNode<T> tree) {
        if (tree == null) return;
        nodeList.add(tree);
        preorderTraverseRecursive(tree.left);
        preorderTraverseRecursive(tree.right);
    }

    public List<TreeNode<T>> preorderTraverse() {
        assert hasBuilt;

        nodeList = new ArrayList<>();
        preorderTraverseRecursive(this.treeRoot);
        return nodeList;
    }

    public List<Byte> serializeTree() {
        assert hasBuilt;
        // make sure pre-order traversed before
        if (nodeList == null) {
            preorderTraverse();
        }
        // Serialize tree to a byte array
        List<Byte> arr = new ArrayList<>();
        for (TreeNode<T> n : nodeList) {
            // We only support serialize of TreeNode<Byte>
            assert n.data == null || n.data instanceof Byte;
            if (n.data == null) {
                // For non-leaf node
                arr.add(NON_LEAF_MARK);
            } else {
                // For leaf node
                arr.add(LEAF_MARK);
                arr.add((Byte) (n.data));
            }
        }
        return arr;
    }

    private TreeNode<T> buildPtr = null;

    private int traverseIndex = 0;

    /**
     * Build tree according to a pre-order traverse arr
     *
     * @param preorderTraverse pre-order traverse arr
     * @return the tree built with `preorderTraverse` begins with `traverseIndex`
     */
    public TreeNode buildTreeFromPreorderRecursive(List<TreeNode> preorderTraverse) {
        assert !hasBuilt;
        assert treeRoot == null;  // assert empty tree
        // --- Recursion end condition
        TreeNode node = preorderTraverse.get(traverseIndex);
        traverseIndex += 1;
        if (node.isLeaf()) {
            return node;
        }
        // --- Do recursion
        node.left = buildTreeFromPreorderRecursive(preorderTraverse);
        node.right = buildTreeFromPreorderRecursive(preorderTraverse);
        return node;
    }

    public TreeNode<T> getTreeRoot() {
        return treeRoot;
    }

}


